import argparse
import backtrader as bt
from datetime import datetime as dt
from time import time
from multiprocessing import cpu_count, Process
import os

from os.path import join, exists, dirname
if __name__ == '__main__':
    import sys
    sys.path.append(dirname(dirname(__file__)))
from strategies.netEasyRelaying import StrategyClass
from utils.localdata import onebyone, download

params = dict(
        mabodyperiod = 22,     # 柱体均值区间
        bodylimit = 1,         # 中柱界定，柱高与柱体均值的比值，0.95代表比均值略小也算中柱
        crossperiod = 3 ,      # 死叉金叉间隔
        planstop = 1.1,        # 计划止盈
        trailpercent = 0.03,   # 跟踪止损
        buyprice = 1,          # 大于0时以形态日收盘价的buyprice倍做限价买入，否则以市价买入
        validdays = 5,         # 限价单的有效期
)

ps = dict(
        bodylimit = [x/100 for x in range(70, 150, 2)],       # 中柱界定，柱高与柱体均值的比值，0.95代表比均值略小也算中柱
        crossperiod = [x for x in range(3, 30, 2)],           # 死叉金叉间隔
)

psv = dict(
        planstop = [x/100 for x in range(105, 131, 3)],     # 计划止盈
        trailpercent = [x/100 for x in range(5, 20, 3)],    # 跟踪止损
        buyprice = [x/100 for x in range(70, 111, 4)],      # 大于0时以形态日收盘价的buyprice倍做限价买入，否则以市价买入
        validdays = [x for x in range(2,11,2)],              # 限价单的有效期
)

ps = dict(
        bodylimit = [0.95,1,1.05],         # 中柱界定，柱高与柱体均值的比值，0.95代表比均值略小也算中柱
        crossperiod = [2,3,5,8,13] ,      # 死叉金叉间隔
        planstop = [1.05,1.06,1.07,1.09,1.12,1.17,1.25,1.48],        # 计划止盈
        trailpercent = [0.03, 0.05,0.07,0.09,0.11],   # 跟踪止损
)

ps2 = None

def pps(p={}, ks=list(ps.keys()), ps=ps, ps2=ps2):
    if len(ks) > 0:
        for x in ps[ks[0]]:
            yield from pps(dict({f'{ks[0]}': x}, **p), ks[1:], ps, ps2)
    elif 'afterdown' in p and p['afterdown'] and ps2:
        yield from pps(p, list(ps2.keys()), ps2, None)
    else:
        yield p

print(len([x for x in pps()]))

def main(logpath, num, pss, fix={}, StrategyClass=StrategyClass):
    logfile = join(logpath, f's{num:03d}_{dt.now().isoformat()}_{os.getpid()}.log')
    def log(text, write=False, file=logfile, cache=dict(a=[])):
        lines = cache['a']
        lines.append(f'{dt.now().isoformat()} {text}')
        if write:
            with open(file, 'a', encoding='utf8') as f:
                f.write('\n'.join(lines))
            cache['a'] = ['\n']

    def logbuy(code, name, date, shape, **kargs):
        log(f'{code} {name} {date} {shape} {kargs["log"]}')

    commparams = dict(log=False, buyhook=logbuy, **fix)
    import re
    replace = re.compile(r'0{4}0+1$')
    for p in pss:
        params = dict(commparams, **p)
        total = [0, 0, 0, 0, 0, 0] # 有交易的证券数，交易总笔数，盈利总笔数，损失总笔数, 证券总数，盈利总和
        # for daily_data, code, name in [d for d in onebyone(klen=300) if d[1] in ['sh.600378', 'sh.603333']]:
        for daily_data, code, name in onebyone(klen=300):
            params['code'] = code
            params['name'] = name
            cerebro = bt.Cerebro()
            # 设置启动资金
            cerebro.broker.setcash(10000000.0)
            # 设置佣金
            cerebro.broker.setcommission(commission=0.001)
            cerebro.addstrategy(StrategyClass, **params)
            data = bt.feeds.PandasData(dataname=daily_data)
            cerebro.adddata(data)
            cerebro.addsizer(bt.sizers.FixedSize, stake=100)  # 每笔交易使用固定交易量，stake代表交易量
            cerebro.addanalyzer(bt.analyzers.TradeAnalyzer, _name='ta')
            result = cerebro.run(plot=False)
            r = result[0].analyzers.ta.get_analysis()
            if r.total and r.total.total > 0:
                s = r.total.total
                l, w, x = 0, 0, cerebro.broker.getvalue() - 10000000
                if r.total.total - r.total.open > 0:
                    l, w = r.lost.total, r.won.total
                log(f'{code} {name} {s} {w} {l} {x:.2f} ')
                total = [a+b for a,b in zip(total, [1,s,w,l,1,x])]
            else:
                total[-2] += 1
        total[-1] = round(total[-1], 2)
        log(replace.sub('', f'={total}={p}={dt.now()}'), True)

def parse_args():
    parser = argparse.ArgumentParser(
        description='行情中继策略参数搜索-多进程')

    parser.add_argument('-n', default=cpu_count()-1, required=False, type=int, help='进程数')

    parser.add_argument('-t', default=False, action='store_const', const=True, help='测交易参数')

    parser.add_argument('-d', dest='download', action='store_const',
                        const=True, default=False, help='是否下载数据')

    parser.add_argument('-o', dest='savepath', default='logs', required=False, help='日志文件存储路径')

    return parser.parse_args()

if __name__ == '__main__':
    args = parse_args()

    if not exists(args.savepath):
        os.makedirs(args.savepath)

    if args.download:
        download()

    title = f"行情中继策略{'交易' if args.t else '选股'}参数搜索"

    params = [x for x in pps(ps=psv, ks=list(psv.keys()))] if args.t else [x for x in pps()]
    size = (len(params) // (args.n - 1)) if args.n > 1 else len(params)
    fix = {'crossperiod':27,'bodylimit':0.85} if args.t else {'validdays':10,'buyprice':0,'trailpercent':0.03,'planstop':1.1}

    for i in range(args.n):
        Process(target=main,kwargs={"logpath":args.savepath, "num": i, "fix": fix, "pss": params[size * i:size * (i + 1)]}).start()

    print(dt.now().isoformat(), f"{title}, 主进程ID:{os.getpid()} 进程数:{args.n} 参数组数量:{len(params)} 每个进程搜索参数组数量：{size}")
